import numpy as np
import os

class TDS:
    """
    TDS类用于计算理论频散谱
    """
    
    def __init__(self, DEBUG=False):
        """
        初始化TDS计算类，设置默认参数
        """
        # 初始化参数
        self.DEBUG = DEBUG
        self.n_layers = 4
        
        # 媒质参数 (与grt.dat中的默认值一致)
        self.media = {
            'z': np.array([0.0, 18.0, 24.0, 30.0, 35.0]),  # 界面深度
            'vp0': np.array([6.0, 6.3, 6.7, 8.2]),  # P波速度
            'vs0': np.array([2.8, 2.92, 3.12, 3.76]),  # S波速度
            'rho': np.array([2.8, 2.9, 3.1, 3.3]),  # 密度
            'qs': np.array([500.0, 500.0, 500.0, 500.0]),  # S波品质因子
            'qp': np.array([1000.0, 1000.0, 1000.0, 1000.0])  # P波品质因子
        }
        
        # 源参数
        self.source = {
            'zs': 10.0,  # 源深度
            'z0': 0.0,  # 接收器深度
            'x0': 0.0,   # 接收器x坐标
            'y0': 1.0,   # 接收器y坐标
            'xs': 0.0,   # 源x坐标
            'ys': 0.0,   # 源y坐标
            'fc': 0.2,   # 中心频率
            'tou': 0.7,  # 上升时间
            'M0': 1e20,  # 标量矩
            'dip': 90.0, # 倾角
            'rake': 45.0,# 滑动角
            'str': 30.0, # 走向
            'mino': 0.01, # 最小频率
            'maxo': 1.0, # 最大频率
            'no': 1000,   # 频率点数
            'minv': 1.0, # 最小速度
            'maxv': 9.0, # 最大速度
            'nv': 1000    # 速度点数
        }
        
        # 其他参数
        self.other_params = {
            'SourceType': 'E',
            'Case': 'A2',
            'Type_STF': 'Bouchon',
            'Outname': 'A2',
            'OutFormat': 'A2',
            'WinSwitch': 'OFF',
            'WinType': 'Hamming',
            'f1': 0.01,
            'f2': 0.07,
            'f3': 96.5,
            'f4': 99.3
        }
        
        # 计算参数
        self.comp = {
            'm': 11,      # 采样参数
            'Twin': 3000.0, # 时间窗口
            'oi': 0.0      # 虚频 (将在初始化时计算)
        }
        
        # 尝试导入Fortran模块
        try:
            import tds4py
            self.fortran_module = tds4py.tds4py
            self.use_fortran = True
        except ImportError:
            self.fortran_module = None
            self.use_fortran = False
            print("警告: 无法导入Fortran模块，将使用纯Python实现(如可用)或无法进行计算")
        
        # 初始化参数
        if self.use_fortran:
            self.fortran_module.init_tds_parameters()
    
    def set_media_parameters(self, z, vp0, vs0, rho, qs, qp):
        """
        设置媒质参数
        
        Args:
            z (numpy.ndarray): 界面深度数组
            vp0 (numpy.ndarray): P波速度数组
            vs0 (numpy.ndarray): S波速度数组
            rho (numpy.ndarray): 密度数组
            qs (numpy.ndarray): S波品质因子数组
            qp (numpy.ndarray): P波品质因子数组
        """
        self.n_layers = len(vp0)
        # self.media['z'] = np.concatenate([z, [z[-1] + 5.0]])  # 添加一个虚拟界面
        self.media['z'] = z  # 添加一个虚拟界面
        self.media['vp0'] = vp0
        self.media['vs0'] = vs0
        self.media['rho'] = rho
        self.media['qs'] = qs
        self.media['qp'] = qp
    
    def set_source_parameters(self, **kwargs):
        """
        设置源参数
        
        Args:
            zs (float): 源深度
            z0 (float): 接收器深度
            x0 (float): 接收器x坐标
            y0 (float): 接收器y坐标
            xs (float): 源x坐标
            ys (float): 源y坐标
            fc (float): 中心频率
            tou (float): 上升时间
            M0 (float): 标量矩
            dip (float): 倾角
            rake (float): 滑动角
            str (float): 走向
            mino (float): 最小频率
            maxo (float): 最大频率
            no (int): 频率点数
            minv (float): 最小速度
            maxv (float): 最大速度
            nv (int): 速度点数
        """
        if self.DEBUG:
            print("TDS:")
            print(kwargs)
        self.source.update(kwargs)
    
    def set_computation_parameters(self, m, Twin):
        """
        设置计算参数
        
        Args:
            m (int): 采样参数
            Twin (float): 时间窗口
        """
        self.comp['m'] = m
        self.comp['Twin'] = Twin
        
        # 重新计算虚频
        dt = Twin / (2**(m-1) * 2 - 1)
        self.comp['oi'] = np.pi / Twin / 10
        
    def write_all_parameters_in_fortran(self):
        """
        设置所有参数到Fortran模块
        """
        if not self.use_fortran:
            return
        
        # "源深度不能超过界面深度"
        assert self.source['zs']<= self.media['z'][-1]

        # 设置媒质参数 - 修复参数传递
        z_vals = np.zeros(self.n_layers + 1, dtype=np.float64)
        vp0_vals = np.zeros(self.n_layers, dtype=np.float64)
        vs0_vals = np.zeros(self.n_layers, dtype=np.float64)
        rho_vals = np.zeros(self.n_layers, dtype=np.float64)
        qs_vals = np.zeros(self.n_layers, dtype=np.float64)
        qp_vals = np.zeros(self.n_layers, dtype=np.float64)
        
        # 正确填充数组
        z_vals[:self.n_layers] = self.media['z'][:self.n_layers]
        z_vals[self.n_layers] = self.media['z'][self.n_layers]  # 最后一个界面深度
        vp0_vals[:] = self.media['vp0'][:self.n_layers]
        vs0_vals[:] = self.media['vs0'][:self.n_layers]
        rho_vals[:] = self.media['rho'][:self.n_layers]
        qs_vals[:] = self.media['qs'][:self.n_layers]
        qp_vals[:] = self.media['qp'][:self.n_layers]
        
        if self.DEBUG:
            print(f"设置媒质参数: n_layers={self.n_layers}")
            print(f"z_vals: {z_vals}")
            print(f"vp0_vals: {vp0_vals}")
            print(f"vs0_vals: {vs0_vals}")
        
        self.fortran_module.set_media_parameters(
            n_layers=int(self.n_layers),  # 确保是整数类型
            z=z_vals,
            vp0=vp0_vals,
            vs0=vs0_vals,
            rho=rho_vals,
            qs=qs_vals,
            qp=qp_vals
        )
        
        # 设置源参数
        self.fortran_module.set_source_parameters(
            float(self.source['zs']), float(self.source['z0']),
            float(self.source['x0']), float(self.source['y0']),
            float(self.source['xs']), float(self.source['ys']),
            float(self.source['fc']), float(self.source['tou']),
            float(self.source['M0']), float(self.source['dip']),
            float(self.source['rake']), float(self.source['str']),
            float(self.source['mino']), float(self.source['maxo']),
            int(self.source['no']), float(self.source['minv']),
            float(self.source['maxv']), int(self.source['nv'])
        )
        
        # 设置计算参数
        self.fortran_module.set_computation_parameters(
            int(self.comp['m']), float(self.comp['Twin']), float(self.comp['oi'])
        )
        
        # 设置其他参数
        self.fortran_module.set_other_parameters(
            self.other_params['SourceType'],
            self.other_params['Case'],
            self.other_params['Type_STF'],
            self.other_params['Outname'],
            self.other_params['OutFormat'],
            self.other_params['WinSwitch'],
            self.other_params['WinType'],
            float(self.other_params['f1']),
            float(self.other_params['f2']),
            float(self.other_params['f3']),
            float(self.other_params['f4'])
        )
        
    def set_other_parameters(self, SourceType, Case, Type_STF, Outname, OutFormat, 
                             WinSwitch, WinType, f1, f2, f3, f4):
        """
        设置其他参数
        
        Args:
            SourceType (str): 源类型
            Case (str): 情况标识
            Type_STF (str): STF类型
            Outname (str): 输出文件名
            OutFormat (str): 输出格式
            WinSwitch (str): 窗口开关
            WinType (str): 窗口类型
            f1 (float): 频率参数1
            f2 (float): 频率参数2
            f3 (float): 频率参数3
            f4 (float): 频率参数4
        """
        self.other_params.update({
            'SourceType': SourceType,
            'Case': Case,
            'Type_STF': Type_STF,
            'Outname': Outname,
            'OutFormat': OutFormat,
            'WinSwitch': WinSwitch,
            'WinType': WinType,
            'f1': f1,
            'f2': f2,
            'f3': f3,
            'f4': f4
        })
    
    def save_parameters(self, filename='grt_generated.dat'):
        """
        将参数保存到文件，用于与原始main.f90结果对比验证
        
        Args:
            filename (str): 保存的文件名
        """
        if not self.use_fortran:
            raise RuntimeError("Fortran模块不可用，无法保存参数")
        
        # 由于我们使用的是简化接口，这里我们手动创建文件
        with open(filename, 'w') as f:
            # 写入接收器参数
            f.write('FOR : receiver information: rarray case:\n\n')
            f.write(' 0.00\n')
            f.write(' 1.00\n')
            f.write(f' {self.source["z0"]:.2f}\n\n')
            
            # 写入源参数
            f.write('FOR: input dat for double-couple source\n\n')
            f.write('100000000000000000000.00\n')
            f.write(f'{self.source["mino"]:.4f}  {self.source["maxo"]:.2f}\n')
            f.write(f' {self.source["no"]}\n')
            f.write(f'{self.source["minv"]:.2f}  {self.source["maxv"]:.2f}\n')
            f.write(f' {self.source["nv"]}\n')
            f.write('45.00 90.00 30.00\n')
            f.write('0.00\n')
            f.write(' 0.00\n')
            f.write(f' {self.source["zs"]:.2f}\n\n')
            
            # 写入计算参数
            f.write('FOR: basic control parameters\n\n')
            f.write('(1). Basic control parameters\n')
            f.write(f'{self.comp["m"]}\n')
            f.write(f'{self.comp["Twin"]:.6f}\n\n')
            
            # 写入源类型
            f.write('(2).Source Typer\n')
            f.write('E\n\n')
            
            # 写入源谱参数
            f.write('(3)Source spectral parameters\n')
            f.write('Bouchon\n\n')
            f.write(' 0.70\n')
            f.write(' 0.00\n')
            f.write(' 0.20\n\n')
            
            # 写入窗口参数
            f.write('(4).Window\n')
            f.write('OFF\n')
            f.write('Hamming\n')
            f.write('0.01\n')
            f.write('0.07\n')
            f.write('96.5\n')
            f.write('99.3\n\n')
            
            # 写入输出参数
            f.write('(5).Out Format\n')
            f.write('A2\n\n')
            
            # 写入层参数
            f.write('FOR: media\n')
            f.write('Number of total layers\n')
            f.write(f'{self.n_layers}\n')
            f.write('Layer ...\n')
            
            for i in range(self.n_layers):
                f.write(f'{i+1}\t {self.media["z"][i]:.4f}\t {self.media["rho"][i]:.4f}\t {self.media["vs0"][i]:.4f}\t {self.media["vp0"][i]:.4f}\t{self.media["qs"][i]:.0f}\t{self.media["qp"][i]:.0f}\n')
        
        print(f"参数已保存到 {filename}")
    
    def calculate_dispersion_spectrum(self):
        """
        计算频散谱（使用频率和速度范围）
        
        Returns:
            numpy.ndarray: 频散谱结果矩阵 (complex)
        """
        if not self.use_fortran:
            raise RuntimeError("Fortran模块不可用，无法计算频散谱")
        
        # 创建结果数组
        result = np.zeros((self.source['no'], self.source['nv']), dtype=np.complex128)
        
        # 调用Fortran函数计算
        result=self.fortran_module.calculate_dispersion_spectrum(
            mino=self.source['mino'], 
            maxo=self.source['maxo'],
            minv=self.source['minv'], 
            maxv=self.source['maxv'],
            no=self.source['no'], 
            nv=self.source['nv'],
        )
        
        return result
    
    def calculate_dispersion_spectrum_arrays(self, freq_array, vel_array):
        """
        计算频散谱（使用频率和速度数组）
        
        Args:
            freq_array (numpy.ndarray): 频率数组
            vel_array (numpy.ndarray): 速度数组
            
        Returns:
            numpy.ndarray: 频散谱结果矩阵 (complex)
        """
        if not self.use_fortran:
            raise RuntimeError("Fortran模块不可用，无法计算频散谱")
        
        # 检查数组长度
        if len(freq_array) != len(vel_array):
            raise ValueError("频率数组和速度数组长度必须相同")
        
        npoints = len(freq_array)
        
        # 创建结果数组
        result = np.zeros(npoints, dtype=np.complex128)
        
        # 调用Fortran函数计算
        result = self.fortran_module.calculate_dispersion_spectrum_arrays(
            freq_array=freq_array, 
            vel_array=vel_array, 
            np=npoints
        )
        result = np.copy(result, order='C')
        
        return result
    
# 使用示例
if __name__ == "__main__":
    # 创建TDS实例
    tds = TDS()
    
    tds._set_all_parameters()

    # 保存参数到文件，用于验证
    tds.save_parameters("grt_generated.dat")
    
    # 计算频散谱
    print("开始计算频散谱...")
    spectrum = tds.calculate_dispersion_spectrum()
    print(f"频散谱计算完成，结果形状: {spectrum.shape}")
    
    # 保存结果
    np.save("dispersion_spectrum.npy", spectrum)
    print("结果已保存到 dispersion_spectrum.npy")
    
    # 使用自定义频率和速度数组计算
    freq_array = np.linspace(0.1, 1.0, 50)
    vel_array = np.linspace(5.0, 9.0, 50)
    
    print("开始计算自定义点频散谱...")
    spectrum_custom = tds.calculate_dispersion_spectrum_arrays(freq_array, vel_array)
    print(f"自定义点频散谱计算完成，结果形状: {spectrum_custom.shape}")
    
    # 保存结果
    np.save("dispersion_spectrum_custom.npy", spectrum_custom)
    print("自定义点结果已保存到 dispersion_spectrum_custom.npy")